# 预计：生成两个气泡重叠的照片
import argparse
import math
import array
import cv2
import numpy as np
import scipy
import os
import random
from matplotlib import pyplot as plt
from scipy.signal import savgol_filter
import ellipsefit_ltl

def img2binary (lr_img):
    '''

    :param lr_img:
    :return: binary, area
    '''
    # lr_img = cv2.imread(img_path)
    [height, weidth, deep] = lr_img.shape
    if weidth < 15 or height < 15:  # 动态的上采样倍数吧
        times = 3
    elif weidth < 30 or height < 30:
        times = 2
    else:
        times = 2

    # 基于三次插值的图像重建
    # times = 2; # 上采样的倍数
    hr_img = cv2.resize(lr_img, (0, 0), fx=times, fy=times, interpolation=cv2.INTER_CUBIC)
    # cv2.imshow('hr_img',hr_img)

    # 局部阈值（动态阈值，平均法）
    gray = cv2.cvtColor(hr_img, cv2.COLOR_BGR2GRAY)
    binary = cv2.adaptiveThreshold(gray, 255, cv2.ADAPTIVE_THRESH_MEAN_C, cv2.THRESH_BINARY, 101, 5)  # 自适应阈值,平均法
    binary = cv2.bitwise_not(binary) # 二值化反转
    # cv2.imshow('binary',binary)

    # 开运算 先腐蚀后膨胀,去除毛刺和小粘连
    kernel = np.ones((5, 5), dtype=np.uint8)
    open = cv2.morphologyEx(binary, cv2.MORPH_OPEN, kernel, 1)
    # cv2.imshow('open', open)

    # 删除连通域面积小于原图5%的部分
    del_small = ellipsefit_ltl.delete_smallarea(open, 0.03)

    # 图像填充,首先填充闭合轮廓
    contours, hierarchy = cv2.findContours(del_small, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE) # cv2.RETR_EXTERNAL是只检测外轮廓
    len_contour = len(contours)
    contour_list = []
    for i in range(len_contour):
        drawing = np.zeros_like(del_small, np.uint8)  # create a black image
        img_contour = cv2.drawContours(drawing, contours, i, (255, 255, 255), -1)
        contour_list.append(img_contour)
    hole_fill = sum(contour_list)

    # 图像填充，再将图像反色，去除小面积的连通域，再反色
    hole_fill = ellipsefit_ltl.inverse_color(hole_fill)
    hole_fill = ellipsefit_ltl.delete_smallarea(hole_fill, 0.1)
    hole_fill = ellipsefit_ltl.inverse_color(hole_fill)

    binary = cv2.resize(hole_fill, (0, 0), fx=1/times, fy=1/times, interpolation=cv2.INTER_CUBIC) # 再将图片下采样到原尺寸

    real_area = 0 # 单个气泡的面积真值
    for a in range(height):
        for b in range(weidth):
            if binary[a, b].all() > 0:
                real_area = real_area + 1

    return binary, real_area

def gen_overlop(num):
    i=0
    while (i<num):
# 随机选取两张图片，并进行二值化
        file_dir = 'D:\\whitebubble\\machinelearning\\yolo\\yolov5-master\\data\\bubbles\\' # 包含677个小气泡
        pathDir = os.listdir(file_dir)  # 取图片的原始路径
        img1 = random.sample(pathDir, 1)  # 随机选取一张图片img1
        img1 = cv2.imread(file_dir+'\\'+(','.join(img1)))
        img2 = random.sample(pathDir, 1)  # 随机选取一张图片img2
        img2 = cv2.imread(file_dir+'\\'+(','.join(img2)))
        # img1 = cv2.resize(img1, (0, 0), fx=10, fy=10, interpolation=cv2.INTER_CUBIC) # 得到的图片扩大十倍
        # mg2 = cv2.resize(img2, (0, 0), fx=10, fy=10, interpolation=cv2.INTER_CUBIC) # 得到的图片扩大十倍
        [height1, weidth1, deep] = img1.shape
        [height2, weidth2, deep] = img2.shape

        binary1, area1 = img2binary(img1)  # 图片二值化，area1为面积的真值
        binary2, area2 = img2binary(img2)

        height = max(height1,height2)
        weidth = max(weidth1,weidth2)
        # img_overlop = np.zeros((height, weidth), np.uint8)  # 创建空画布
        img_overlop = np.zeros((height, weidth))  # 创建空画布

        img_overlop[int((height-height1)/2):int((height+height1)/2), int((weidth-weidth1)/2):int((weidth+weidth1)/2)] = binary1 # 第一张图放在中间
        # cv2.imshow('img_overlop1', img_overlop)

        # 为第二张图生成一个随机的中心点
        angle = np.random.uniform(-180,180) # 随机角度
        theta = angle*(math.pi/180.0)
        half_diag = pow(pow(height/2,2)+pow(weidth/2,2),0.5) # 随机距离，范围1/2对角线长
        distance = np.random.uniform(-half_diag,half_diag)

        center2_y = (height/2)+distance*math.cos(theta)
        center2_x = (weidth/2)+distance*math.sin(theta)

        # 第二张图的区域
        point_up = int((center2_y-(height2)/2)//1) # 在overlop图中的坐标
        point_down = int((center2_y+(height2)/2)//1)
        point_left = int((center2_x-(weidth2)/2)//1)
        point_right = int((center2_x+(weidth2)/2)//1)
        if point_up<0 and point_left<0: # 左上角出界
            A = img_overlop[0:point_down, 0:point_right] * binary2[height2 - point_down:height2,weidth2 - point_right:weidth2]
            img_overlop[0:point_down, 0:point_right] =img_overlop[0:point_down, 0:point_right] + binary2[height2-point_down:height2, weidth2-point_right:weidth2]
            # cv2.imshow('sss',img_overlop)
        elif point_up<0 and point_right>weidth:  # 右上角出界
            A = img_overlop[0:point_down, point_left:weidth] * binary2[height2 - point_down:height2, 0:weidth-point_left]
            img_overlop[0:point_down, point_left:weidth] =img_overlop[0:point_down, point_left:weidth] + binary2[height2 - point_down:height2, 0:weidth-point_left]
        elif point_left<0 and point_down>height: # 左下角出界
            A = img_overlop[point_up:height, 0:point_right] * binary2[0:height - point_up, weidth2 - point_right:weidth2]
            img_overlop[point_up:height, 0:point_right] =img_overlop[point_up:height, 0:point_right] + binary2[0:height-point_up, weidth2-point_right:weidth2]
        elif point_down>height and point_right>weidth: # 右下角出界
            A = img_overlop[point_up:height, point_left:weidth] * binary2[0:height - point_up, 0:weidth-point_left]
            img_overlop[point_up:height, point_left:weidth] =img_overlop[point_up:height, point_left:weidth] + binary2[0:height - point_up, 0:weidth-point_left]
        elif point_up<0: # 上边出界
            A = img_overlop[0:point_down, point_left:point_right] * binary2[height2 - point_down:height2, 0:point_right-point_left]
            img_overlop[0:point_down, point_left:point_right] =img_overlop[0:point_down, point_left:point_right] + binary2[height2 - point_down:height2, 0:point_right-point_left]
        elif point_left<0: # 左边出界
            A = img_overlop[point_up:point_down, 0:point_right] * binary2[0:height2, weidth2 - point_right:weidth2]
            img_overlop[point_up:point_down, 0:point_right] =img_overlop[point_up:point_down, 0:point_right] + binary2[0:height2, weidth2 - point_right:weidth2]
        elif point_right>weidth: # 右边出界
            A = img_overlop[point_up:point_down, point_left:weidth] * binary2[0:height2, 0:weidth-point_left]
            img_overlop[point_up:point_down, point_left:weidth] =img_overlop[point_up:point_down, point_left:weidth] + binary2[0:height2, 0:weidth-point_left]
        elif point_down>height: # 下边出界
            A = img_overlop[point_up:height, point_left:point_right] * binary2[0:height-point_up, 0:weidth2]
            img_overlop[point_up:height, point_left:point_right] =img_overlop[point_up:height, point_left:point_right] + binary2[0:height-point_up, 0:weidth2]
        else:
            A = img_overlop[point_up:point_down, point_left:point_right] * binary2
            img_overlop[point_up:point_down, point_left:point_right] =img_overlop[point_up:point_down, point_left:point_right] + binary2

        img_overlop[img_overlop>255] = 255
        img_overlop = np.array(img_overlop,dtype='uint8')
        overlop_pixel = np.sum(A >= 1)  # 计算叠加两张图片时，重叠的像素数。A是两张图片对应位置相乘，非零即重叠。条件为大于等于1
        overlop_ratio = overlop_pixel/area1
        # img_overlop = cv2.resize(img_overlop, (0, 0), fx=10, fy=10, interpolation=cv2.INTER_CUBIC) # 得到的图片扩大十倍
        cv2.imwrite('gen_overlop0831\\%d.jpg'%i,img_overlop)
        file = open('gen_overlop0831\\%d.txt'%i,'w')
        file.write(str(area1)+'\n'+str(overlop_ratio)) # 第一行为中间气泡的面积真值，第二行为重叠比例
        file.close()

        i+=1
        print(i)
        # return img_overlop, area1

gen_overlop(1000)